namespace KillerKitty {


	//The detection class
	//this is the class actually responsible for detecting the virus.

	class kkDetectionEngine : public KillerKittyBase{
	private:
		long sigCount;
		unsigned char *hexChars;
		kkSignatureParser *parser;
		unsigned char *fileBuffer;

		long getStartOffset(char *off, unsigned char *file, long fsize, bool *star);
		inline unsigned char getCharValue(unsigned char *buf, long index, int useLowerNibble) {
			unsigned char val = 0;
			if(useLowerNibble) 
				val= hexChars[(buf[index]&0x0F)];
			else 
				val= hexChars[(buf[index]&0xF0)>>4];	//FIXME: add endian-ness checking? is it endian safe?

			return val;
		}

	public:
		kkDetectionEngine();
		~kkDetectionEngine();
		bool initEngine();

		//Accepts a filename and returns whether a virus was detected or not
		string isMalicious(char* fname);

	};

	kkDetectionEngine::kkDetectionEngine() {

		//debug
		log.write(this,"Detection Engine init Successful",debug);
		parser = new kkClamSignatureParser();
		sigCount = 0;
		fileBuffer = NULL;
		hexChars = (unsigned char *) "0123456789abcdef";

	}

	kkDetectionEngine::~kkDetectionEngine() {
		delete parser;
	}

	bool kkDetectionEngine::initEngine() {
		sigCount=parser->parseSignatures();

		if(sigCount<=0) {
			log.write(this,"Error parsing signatures!",error);
			return false;
		}
		else {
			char msg[20];
			ltoa(sigCount,msg,10);
			string s("Succesfully parsed ");
			s+=msg;
			s+=" signatures";
			log.write(this,s,info);
		}

		return true;
	}

	string kkDetectionEngine::isMalicious(char *fname) {

		//long lTotalRegexCnt=0, lRegExesPerSig=0;
		//long lTotalCount=0, lMaxRegexes=0, lOffsetStar=0;

		ifstream file(fname,ios::binary);
		string sFile(fname);
		kkSignatureParser& p = *parser;
		long lfileSize=-1;

		
		if(!file.is_open()) {	//if file cannot be read
			log.write(this,"Cannot open specified file : "+sFile,error);
			return "-1";
		}
		else {	//open file and copy contents to buffer
			log.write(this,"Opened "+sFile,info);
			file.seekg(0, ios_base::end );
			lfileSize = file.tellg();
			file.seekg(0, ios_base::beg );

			fileBuffer = new unsigned char[lfileSize+1];
			file.read((char *)fileBuffer,lfileSize);
			fileBuffer[lfileSize] = '\0';

			//cout<<"\nFile size : "<<lfileSize;
			//cout<<"\nContents : "<<(char *)fileBuffer<<endl;
		}

		char digest[33];  //32 for md5 hash + 1 fr null term

		//iterate through each of the available signatures and check if a match is found
		//TODO:OPTIMIZE! -- what follows are possibly the most time consuming loops of the whole program
		//ideas:		
		//for dynamic:
		//...?		
			
		//Performing Binary Search on static signatures.

		pair< unsigned long, unsigned long > indices(p.getSigRange(sigStatic));
		unsigned long low=indices.first,high=indices.second,mid;		
		while(low<=high){
			mid=(low+high)/2;		

			if(lfileSize < p[mid].fileLength)
				high=mid-1;
			else if(lfileSize > p[mid].fileLength)
				low=mid+1;
			else{
				//File length matches ! Generating MD5
				MD5::generateMD5(fname,digest);
				log.write(this,"MD5 is : "+string(digest)+"",info);
				//Looping to check for each signature that has the same file length
				for(unsigned long ind=mid ; ind < p.getSigCount() && p[ind].fileLength==lfileSize ; ind++){
					
					if(stricmp(p[ind].hexcode,digest)==0) {	//if md5 in signature info struct matches md5 generated for the file, we have found a virus
						string name(p[ind].name);
						log.write(this,"Detected "+name+"!",info);
						cout<<"Detected :";
						p.printDetails(p[ind]);
						delete fileBuffer;
						return name;
					}
				}
				
				break;
			}               
		}//binary search ends
		
		pair< unsigned long, unsigned long > dynInd(p.getSigRange(sigDynamic));

		for(long i=dynInd.first;i<=dynInd.second;i++) {
			kkSignature_t& sig = p[i];

			bool offsetStar = false;
			long lStartOffset = getStartOffset(sig.offset, fileBuffer, lfileSize, &offsetStar);
			long lBufIndex=0,lSigIndex=0;	//note : sig index should be incremented by 2 for every 1 increment to lBuffIndex
			long lSigLength = strlen(sig.hexcode);
			if(sig.offset[0] == '*' ) offsetStar = true;

			//Presence of a malware is not posssible if the offset becomes negative.
			if(lStartOffset < 0 && (lStartOffset + lfileSize) < 0) {				
				return "0";
				delete fileBuffer;
			}		
			
			if(sig.sigType==sigDynamic) {
				//FIXME: refactor!
				if(!sig.isRegExPresent(sig.hexcode)) {
					//do string search through the file for the sig.
					//TODO: Replace with Booyer-Moore or better
					bool done = false;
					bool match = false;
					while(!done) {	

						if(lBufIndex > lSigLength || (lBufIndex+lStartOffset) > lfileSize || lSigLength/2 > lfileSize ) { //if end of sig or eof is reached
							done = true;	
							continue;
						}

						//if higher nibble in file matches higher nibble in sig
						if(sig.hexcode[lSigIndex] != getCharValue(fileBuffer, lStartOffset + lBufIndex, lSigIndex % 2 )) {
							//cout<<endl<<sig.hexcode[lSigIndex]<<" in sig does not match "<<getCharValue(fileBuffer, lStartOffset + lBufIndex, lSigIndex % 2 );
							done = true;
							continue;
						}
						lSigIndex++;
						//if lower nibble in file matches lower nibble in sig
						if(sig.hexcode[lSigIndex] != getCharValue(fileBuffer, lStartOffset + lBufIndex, lSigIndex % 2)) {
							//cout<<endl<<sig.hexcode[lSigIndex]<<" in sig does not match "<<getCharValue(fileBuffer, lStartOffset + lBufIndex, lSigIndex % 2 );
							done = true;
							continue;
						}
						lBufIndex++;
						lSigIndex++;

						if(lSigIndex == lSigLength) { //it is a match only if the end of the signature has been reached.
							done=true;
							match = true;	
							break;
						}
					}//end while

					if(match) {	//virus matched
						string name(sig.name);
						log.write(this,"Detected "+name+"!",info);
						cout<<"Detected :";
						p.printDetails(sig);
						delete fileBuffer;
						return name;
					}
				}
				else {
					kkRegExEngine r;
					
					bool match = false;

					//forward search until a regex char is found, or sig mismatch
					while( lStartOffset < lfileSize ) {	
						
						for(lBufIndex = lSigIndex = 0; lStartOffset+lBufIndex<lfileSize; lBufIndex++ ) {
							//check higher nibble
							if(sig.hexcode[lSigIndex] != getCharValue(fileBuffer, lStartOffset + lBufIndex, lSigIndex % 2 )) {
								if(sig.isCharRegEx(sig.hexcode[lSigIndex])) match=true;
								break;
							}

							lSigIndex++;
							//if lower nibble in file matches lower nibble in sig
							if(sig.hexcode[lSigIndex] != getCharValue(fileBuffer, lStartOffset + lBufIndex, lSigIndex % 2)) {
								if(sig.isCharRegEx(sig.hexcode[lSigIndex])) match=true;
								break;
							}
							lSigIndex++;
						}

						if(offsetStar && !match) lStartOffset++;	//search from next char for possible match
						else break;
					}
						
					if( !match ) continue; //No need for reverse search if forward search returns nothing

					lBufIndex = lfileSize-1;
					match = false;
					long lEndOffset = lfileSize - 1;
					//reverse search until a regex char is found, or sig mismatch
					while( lEndOffset > lStartOffset ) {

						for(lSigIndex = lSigLength -1, lBufIndex = 0; lSigIndex > 0; lBufIndex++ ) {

							//if lower nibble in file matches lower nibble in sig
							if(sig.hexcode[lSigIndex] != getCharValue(fileBuffer, lEndOffset - lBufIndex, lSigIndex % 2)) {
								if(sig.isCharRegEx(sig.hexcode[lSigIndex])) match = true;
								break;
							}

							lSigIndex--;
							//higher nibble
							if(sig.hexcode[lSigIndex] != getCharValue(fileBuffer, lEndOffset - lBufIndex, lSigIndex % 2 )) {
								if(!sig.isCharRegEx(sig.hexcode[lSigIndex])) match = true;
								break;
							}
							lSigIndex--;
						}

						if(match) break;
						else lEndOffset--;
					}//end while

					//possible match
					if(match) {
						char *hex = new char[(lEndOffset-lStartOffset+1)*2+1]; //*2 -> actual length, +1 more for \0
						
						//UGLY! :(( FIXME: REFACTOR!
						long j=0;
						cout<<"\nlStartOffset : "<<lStartOffset<<"  lEndOffset  : "<<lEndOffset<<endl; 
						for(long i=lStartOffset*2;i<(lEndOffset+1)*2;i++,j++) hex[j] = getCharValue(fileBuffer, i/2, i % 2 );
						hex[j] = '\0';

						//cout<<"Possible match "<<sig.name<<"\nFrom : "<<lStartOffset<<" to "<<lEndOffset<<endl;
						
						//if virus sig matched file contents
						
						int result = r.isMatch(sig.hexcode,hex,(lEndOffset-lStartOffset+1)*2);
						if( result==1 ) {
							string name(sig.name);
							log.write(this,"Detected "+name+"!",info);
							cout<<"Detected :";
							p.printDetails(sig);
							delete fileBuffer;
							delete hex;
							return name;
						}
						else if( result == -1) {
							//TODO: Write regex engine.. :((, or check based on file size
							return "-1";
						}
						delete hex;
					}

				}
			}
			else {	//if sigType is anything other than static or dynamic
				log.write(this, "Unsupported sigType!",error);
				cout<<"Unsupported sigType!"<<endl;
				p.printDetails(sig);

			}
		}
		//log.write(this,"No virus detected!",info);
		delete fileBuffer;
		return "0";
	}

	long kkDetectionEngine::getStartOffset(char *off, unsigned char *file, long fsize, bool *star) {
		long offset = 0;

		//TODO: if determining seg start or EP fails, then set offset * as backup
		switch(*off) {
			case 's':
			case 'S':
				//segment
				//TODO: determine correct seg start and set offset
				offset=PEReader::getSectionAddress(file, atoi(off+1));
				if(offset > fsize) {
					offset = 0;	//FIXME: track down cause and fix
					*star = true;
				}
				break;

			//eof or entry point
			case 'e':
			case 'E':
				
				if(*(off+1) == 'o' || *(off+1) == 'O') offset =-atol(off+4);	//if eof
				//find EP
				else {
					offset = PEReader::getEP(file);
					if(offset>fsize) {
						offset = 0;	//FIXME: track down cause and fix
						*star = true;
					}
				}
				break;

			//none and *
			case 'n':
			case 'N':
			case '*':
				offset = 0;
				break;
			//absolute offsets
			default :
				offset = atol(off);

		}
		if(strcmp(off,"EOF")==0){
			cout<<"\n\nOffset String : "<<off<<"\nReturn value : "<<offset;
			cin.get();
		}
		return offset;	//if <0, its wrt EOF, if >0, it is absolute, if ==0 then any offset
	}

};
